

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.ops.operations.comm_ops &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.ops.operations.comm_ops</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.ops.operations.comm_ops</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2020 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>

<span class="sd">&quot;&quot;&quot;Communication APIs.</span>
<span class="sd">&quot;&quot;&quot;</span>

<span class="kn">from</span> <span class="nn">mindspore.common</span> <span class="kn">import</span> <span class="n">Tensor</span>
<span class="kn">from</span> <span class="nn">..._checkparam</span> <span class="kn">import</span> <span class="n">Validator</span> <span class="k">as</span> <span class="n">validator</span>
<span class="kn">from</span> <span class="nn">..._checkparam</span> <span class="kn">import</span> <span class="n">Rel</span>
<span class="kn">from</span> <span class="nn">...communication.management</span> <span class="kn">import</span> <span class="n">get_rank</span><span class="p">,</span> <span class="n">get_group_size</span><span class="p">,</span> <span class="n">GlobalComm</span><span class="p">,</span> <span class="n">_get_group</span>
<span class="kn">from</span> <span class="nn">...common</span> <span class="kn">import</span> <span class="n">dtype</span> <span class="k">as</span> <span class="n">mstype</span>
<span class="kn">from</span> <span class="nn">..primitive</span> <span class="kn">import</span> <span class="n">PrimitiveWithInfer</span><span class="p">,</span> <span class="n">PrimitiveWithCheck</span><span class="p">,</span> <span class="n">Primitive</span><span class="p">,</span> <span class="n">prim_attr_register</span>
<span class="kn">from</span> <span class="nn">...common.api</span> <span class="kn">import</span> <span class="n">context</span>


<div class="viewcode-block" id="ReduceOp"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.ReduceOp.html#mindspore.ops.ReduceOp">[docs]</a><span class="k">class</span> <span class="nc">ReduceOp</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Operation options for reducing tensors. This is an enumerated type, not an operator.</span>
<span class="sd">    Mainly used in data parallel mode.</span>

<span class="sd">    The main calling methods are as follows:</span>

<span class="sd">    - SUM: ReduceOp.SUM.</span>
<span class="sd">    - MAX: ReduceOp.MAX.</span>
<span class="sd">    - MIN: ReduceOp.MIN.</span>
<span class="sd">    - PROD: ReduceOp.PROD.</span>

<span class="sd">    There are four kinds of operation options, &quot;SUM&quot;, &quot;MAX&quot;, &quot;MIN&quot;, and &quot;PROD&quot;.</span>

<span class="sd">    - SUM: Take the sum.</span>
<span class="sd">    - MAX: Take the maximum.</span>
<span class="sd">    - MIN: Take the minimum.</span>
<span class="sd">    - PROD: Take the product.</span>

<span class="sd">    Note:</span>
<span class="sd">        For more, refer to example. This needs to run in an environment with multiple graphics cards.</span>
<span class="sd">        The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops import ReduceOp</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.allreduce_sum = ops.AllReduce(ReduceOp.SUM)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         return self.allreduce_sum(x)</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; input_ = Tensor(np.ones([2, 8]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[2. 2. 2. 2. 2. 2. 2. 2.]</span>
<span class="sd">         [2. 2. 2. 2. 2. 2. 2. 2.]]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">SUM</span> <span class="o">=</span> <span class="s2">&quot;sum&quot;</span>
    <span class="n">MAX</span> <span class="o">=</span> <span class="s2">&quot;max&quot;</span>
    <span class="n">MIN</span> <span class="o">=</span> <span class="s2">&quot;min&quot;</span>
    <span class="n">PROD</span> <span class="o">=</span> <span class="s2">&quot;prod&quot;</span></div>


<span class="n">target_dtypes</span> <span class="o">=</span> <span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">check_hcom_group_valid</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Check if hcom group is valid.&quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;mode&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="n">context</span><span class="o">.</span><span class="n">PYNATIVE_MODE</span> <span class="ow">and</span> \
            <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;device_target&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;Ascend&quot;</span> <span class="ow">and</span> \
            <span class="n">group</span> <span class="o">!=</span> <span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;group&#39; only support &#39;hccl_world_group&#39; in pynative mode, but got &quot;</span>
                           <span class="sa">f</span><span class="s2">&quot;&#39;group&#39;: </span><span class="si">{</span><span class="n">group</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">AllReduce</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces the tensor data across all devices in such a way that all devices will get the same final result.</span>

<span class="sd">    Note:</span>
<span class="sd">        The operation of AllReduce does not support &quot;prod&quot; currently.</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection. The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        op (str): Specifies an operation used for element-wise reductions,</span>
<span class="sd">                  like sum, max, and min. Default: ReduceOp.SUM.</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.</span>
<span class="sd">        The contents depend on the specified operation.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If any of `op` and `group` is not a str,</span>
<span class="sd">                   or fusion is not an integer, or the input&#39;s dtype is bool.</span>
<span class="sd">        ValueError: If the `op` is &quot;prod&quot;.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with two devices. Refer to the tutorial &gt; Distributed Training on mindspore.cn</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops import ReduceOp</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.allreduce_sum = ops.AllReduce(ReduceOp.SUM)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         return self.allreduce_sum(x)</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; input_ = Tensor(np.ones([2, 8]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[2. 2. 2. 2. 2. 2. 2. 2.]</span>
<span class="sd">         [2. 2. 2. 2. 2. 2. 2. 2.]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize AllReduce.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">)):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;op&#39; should be str, but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">op</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;group&#39; should be str, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">check_hcom_group_valid</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">op</span> <span class="o">=</span> <span class="n">op</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;fusion&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;index&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">AllGather</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gathers tensors from the specified communication group.</span>

<span class="sd">    Note:</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection. The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor. If the number of devices in the group is N,</span>
<span class="sd">        then the shape of output is :math:`(N, x_1, x_2, ..., x_R)`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `group` is not a str.</span>
<span class="sd">        ValueError: If the local rank id of the calling process in the group</span>
<span class="sd">                    is larger than the group&#39;s rank size.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with two devices. Refer to the tutorial &gt; Distributed Training on mindspore.cn</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, context</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE)</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.allgather = ops.AllGather()</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         return self.allgather(x)</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; input_x = Tensor(np.ones([2, 8]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[1. 1. 1. 1. 1. 1. 1. 1.]</span>
<span class="sd">         [1. 1. 1. 1. 1. 1. 1. 1.]</span>
<span class="sd">         [1. 1. 1. 1. 1. 1. 1. 1.]</span>
<span class="sd">         [1. 1. 1. 1. 1. 1. 1. 1.]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize AllGather.&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">=</span> <span class="n">get_rank</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span> <span class="o">=</span> <span class="n">get_group_size</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s1">&#39;rank&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank</span><span class="p">,</span> <span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">LT</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;fusion&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;mean_flag&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">),</span> <span class="s2">&quot;x shape&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>


<span class="k">class</span> <span class="nc">_MiniStepAllGather</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do reducescatter in backward in mini-step. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group to work on. Default: None.</span>
<span class="sd">        grad_accumulation_step (int): The grad accumulation step. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">,</span> <span class="n">grad_accumulation_step</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">mean_flag</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _MiniStepAllGather.&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">=</span> <span class="n">get_rank</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span> <span class="o">=</span> <span class="n">get_group_size</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s1">&#39;rank&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank</span><span class="p">,</span> <span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">LT</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;fusion&#39;</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">grad_accumulation_step</span> <span class="o">=</span> <span class="n">grad_accumulation_step</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mean_flag</span> <span class="o">=</span> <span class="n">mean_flag</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">),</span> <span class="s2">&quot;x shape&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">_MicroStepAllGather</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do reducescatter in backward in mini-step. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group to work on. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">,</span> <span class="n">mean_flag</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">=</span> <span class="n">get_rank</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span> <span class="o">=</span> <span class="n">get_group_size</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s1">&#39;rank&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank</span><span class="p">,</span> <span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">LT</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;fusion&#39;</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;do_mirror&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mean_flag</span> <span class="o">=</span> <span class="n">mean_flag</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">),</span> <span class="s2">&quot;x shape&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">z_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">_HostAllGather</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gathers tensors from the specified communication group on host.</span>

<span class="sd">    Note:</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection.</span>
<span class="sd">        _HostAllGather is a host-side operator, it depends on OpenMPI and must use build option -M on</span>
<span class="sd">        to enable it. Using mpirun command to run it:</span>
<span class="sd">        mpirun -output-filename log -merge-stderr-to-stdout -np 3 python test_host_all_gather.py</span>

<span class="sd">    Args:</span>
<span class="sd">        group (Union[tuple[int],list[int]]): The rand_ids of communication group to work on. Default: None.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If group is not a list nor tuple, or elements of group are not int.</span>
<span class="sd">        ValueError: If group is not set, or rank_id from group not in [0, 7].</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor. If the number of devices in the group is N,</span>
<span class="sd">        then the shape of output is :math:`(N, x_1, x_2, ..., x_R)`.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _HostAllGather.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">group</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;group&#39; cannot be None, but got </span><span class="si">{</span><span class="n">group</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">group</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="mi">2</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">GE</span><span class="p">,</span> <span class="s2">&quot;group size&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">r</span> <span class="ow">in</span> <span class="n">group</span><span class="p">:</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_int_range</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">7</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">INC_BOTH</span><span class="p">,</span> <span class="s2">&quot;rank_id&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;rank_id&quot;</span><span class="p">,</span> <span class="n">r</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">group_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">group</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">group</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_positive_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">),</span> <span class="s2">&quot;x shape&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">group_size</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>


<div class="viewcode-block" id="ReduceScatter"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.ReduceScatter.html#mindspore.ops.ReduceScatter">[docs]</a><span class="k">class</span> <span class="nc">ReduceScatter</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces and scatters tensors from the specified communication group.</span>

<span class="sd">    Note:</span>
<span class="sd">        The back propagation of the op is not supported yet. Stay tuned for more.</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection. The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        op (str): Specifies an operation used for element-wise reductions,</span>
<span class="sd">                  like SUM, MAX, AVG. Default: ReduceOp.SUM.</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If any of operation and group is not a string.</span>
<span class="sd">        ValueError: If the first dimension of the input cannot be divided by the rank size.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with two devices. Refer to the tutorial &gt; Distributed Training on mindspore.cn</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, context</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops import ReduceOp</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE)</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.reducescatter = ops.ReduceScatter(ReduceOp.SUM)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         return self.reducescatter(x)</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; input_ = Tensor(np.ones([8, 8]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[2. 2. 2. 2. 2. 2. 2. 2.]</span>
<span class="sd">         [2. 2. 2. 2. 2. 2. 2. 2.]</span>
<span class="sd">         [2. 2. 2. 2. 2. 2. 2. 2.]</span>
<span class="sd">         [2. 2. 2. 2. 2. 2. 2. 2.]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize ReduceScatter.&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;op&#39;</span><span class="p">,</span> <span class="n">op</span><span class="p">,</span> <span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">),),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">op</span> <span class="o">=</span> <span class="n">op</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span> <span class="o">=</span> <span class="n">get_group_size</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;rank_size&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;fusion&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;rank_size&#39; cannot be zero, but got </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the first dimension of &#39;x_shape&#39; should be divided by &#39;rank_size&#39;, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;x_shape[0]&#39;: </span><span class="si">{</span><span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2">, &#39;rank_size&#39;: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank_size</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>


<span class="k">class</span> <span class="nc">_HostReduceScatter</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces and scatters tensors from the specified communication group on host.</span>

<span class="sd">    Note:</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection.</span>
<span class="sd">        _HostReduceScatter is a host-side operator, it depends on OpenMPI and must use build option</span>
<span class="sd">        -M on to enable it. Using mpirun command to run it:</span>
<span class="sd">        mpirun -output-filename log -merge-stderr-to-stdout -np 3 python test_host_reduce_scatter.py</span>

<span class="sd">    Args:</span>
<span class="sd">        op (str): Specifies an operation used for element-wise reductions,</span>
<span class="sd">                  like sum, max, avg. Default: ReduceOp.SUM.</span>
<span class="sd">        group (Union[tuple[int],list[int]]): The rand_ids of communication group to work on. Default: None.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If op is not a string and group is not a list nor tuple,</span>
<span class="sd">                   or elements of group are not int.</span>
<span class="sd">        ValueError: If the first dimension of input can not be divided by group size,</span>
<span class="sd">                    or group is not set, or rank_id not in [0, 7].</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _HostReduceScatter.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">group</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;group&#39; cannot be None, but got </span><span class="si">{</span><span class="n">group</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;op&#39;</span><span class="p">,</span> <span class="n">op</span><span class="p">,</span> <span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">),),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">group</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="mi">2</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">GE</span><span class="p">,</span> <span class="s2">&quot;group size&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">r</span> <span class="ow">in</span> <span class="n">group</span><span class="p">:</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_int_range</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">7</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">INC_BOTH</span><span class="p">,</span> <span class="s2">&quot;rank_id&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;rank_id&quot;</span><span class="p">,</span> <span class="n">r</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">op</span> <span class="o">=</span> <span class="n">op</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">group_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">group</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">group</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">group_size</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the first dimension of &#39;x_shape&#39; should be divided by &#39;group_size&#39;, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;x_shape[0]&#39;: </span><span class="si">{</span><span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2">, &#39;rank_size&#39;: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">group_size</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">group_size</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>


<span class="k">class</span> <span class="nc">Broadcast</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Broadcasts the tensor to the whole group.</span>

<span class="sd">    Note:</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection. The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        root_rank (int): Source rank. Required in all processes except the one</span>
<span class="sd">                   that is sending the data.</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.</span>
<span class="sd">        The contents depend on the data of the `root_rank` device.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If root_rank is not an integer or group is not a string.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with multiple processes.</span>
<span class="sd">        &gt;&gt;&gt; # Please refer to the tutorial &gt; Distributed Training on mindspore.cn.</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import context</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE)</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.broadcast = ops.Broadcast(1)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         return self.broadcast((x,))</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; input_x = Tensor(np.ones([2, 4]).astype(np.int32))</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape[2,4], dtype=Int32, value=</span>
<span class="sd">        [[1, 1, 1, 1],</span>
<span class="sd">         [1, 1, 1, 1]]),)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">root_rank</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Broadcast.&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;root_rank&#39;</span><span class="p">,</span> <span class="n">root_rank</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">check_hcom_group_valid</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x_dtype</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;input_x&#39; should be a tuple, but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">x_dtype</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">!&quot;</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">_ele</span> <span class="ow">in</span> <span class="n">x_dtype</span><span class="p">:</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">_ele</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">_AllSwap</span><span class="p">(</span><span class="n">PrimitiveWithCheck</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    _AllSwap is a collective operation.</span>

<span class="sd">    _AllSwap sends data from the all processes to the all processes in the specified group. It has two phases:</span>

<span class="sd">    - The scatter phase: On each process, the operand is split into the send size of blocks along the</span>
<span class="sd">      0-th axis, and the blocks are scattered to all processes, e.g., the ith block is send to the ith process.</span>
<span class="sd">    - The gather phase: Each process concatenates the received blocks along the 0-th axis.</span>

<span class="sd">    Note:</span>
<span class="sd">        The tensors must have the same format in all processes of the collection.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group name.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        tensor_in (tensor): A 2-D tensor. On each process, divide blocks into number of the send size.</span>
<span class="sd">        send_size (tensor): A 1-D int64 tensor. The element is the send data size for each process.</span>
<span class="sd">        recv_size (tensor): A 1-D int64 tensor. The element is the receive data size for each process.</span>

<span class="sd">    Returns:</span>
<span class="sd">        tensor_out (tensor): The result tensor.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If group is not a string.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _AllSwap&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_prim_io_names</span><span class="p">(</span><span class="n">inputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;tensor_in&#39;</span><span class="p">,</span> <span class="s1">&#39;send_size&#39;</span><span class="p">,</span> <span class="s1">&#39;recv_size&#39;</span><span class="p">],</span> <span class="n">outputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;tensor_out&#39;</span><span class="p">])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__check__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor_in</span><span class="p">,</span> <span class="n">send_size</span><span class="p">,</span> <span class="n">recv_size</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_subclass</span><span class="p">(</span><span class="s2">&quot;tensor_in&quot;</span><span class="p">,</span> <span class="n">tensor_in</span><span class="p">[</span><span class="s1">&#39;dtype&#39;</span><span class="p">],</span> <span class="n">mstype</span><span class="o">.</span><span class="n">tensor</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s2">&quot;send_size&quot;</span><span class="p">,</span> <span class="n">send_size</span><span class="p">[</span><span class="s1">&#39;dtype&#39;</span><span class="p">],</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">int64</span><span class="p">],</span>
                                           <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s2">&quot;recv_size&quot;</span><span class="p">,</span> <span class="n">recv_size</span><span class="p">[</span><span class="s1">&#39;dtype&#39;</span><span class="p">],</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">int64</span><span class="p">],</span>
                                           <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

        <span class="n">validator</span><span class="o">.</span><span class="n">check_equal_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">tensor_in</span><span class="p">[</span><span class="s1">&#39;shape&#39;</span><span class="p">]),</span> <span class="mi">2</span><span class="p">,</span> <span class="s2">&quot;tensor_in&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_equal_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">send_size</span><span class="p">[</span><span class="s1">&#39;shape&#39;</span><span class="p">]),</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;send_size&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_equal_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">recv_size</span><span class="p">[</span><span class="s1">&#39;shape&#39;</span><span class="p">]),</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;recv_size&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

        <span class="n">out_shape</span> <span class="o">=</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="p">[</span><span class="n">tensor_in</span><span class="p">[</span><span class="s1">&#39;shape&#39;</span><span class="p">][</span><span class="mi">1</span><span class="p">]]</span>
        <span class="n">out</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;shape&#39;</span><span class="p">:</span> <span class="n">out_shape</span><span class="p">,</span>
               <span class="s1">&#39;dtype&#39;</span><span class="p">:</span> <span class="n">tensor_in</span><span class="p">[</span><span class="s1">&#39;dtype&#39;</span><span class="p">],</span>
               <span class="s1">&#39;value&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">}</span>
        <span class="k">return</span> <span class="n">out</span>


<div class="viewcode-block" id="NeighborExchange"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.NeighborExchange.html#mindspore.ops.NeighborExchange">[docs]</a><span class="k">class</span> <span class="nc">NeighborExchange</span><span class="p">(</span><span class="n">Primitive</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    NeighborExchange is a collective operation.</span>

<span class="sd">    NeighborExchange sends data from the local rank to ranks in the send_rank_ids,</span>
<span class="sd">    as while receive data from recv_rank_ids.</span>

<span class="sd">    Note:</span>
<span class="sd">        The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">        This operator requires a full-mesh network topology, each device has the same vlan id, and the ip &amp; mask are</span>
<span class="sd">        in the same subnet, please check the `details \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/programming_guide/zh-CN/r1.6/distributed_training_ops.html#id2&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        send_rank_ids (list(int)): Ranks which the data is sent to.</span>
<span class="sd">        recv_rank_ids (list(int)): Ranks which the data is received from.</span>
<span class="sd">        recv_shapes (tuple(list(int))): Data shape which received from recv_rank_ids.</span>
<span class="sd">        send_shapes (tuple(list(int))): Data shape which send to the send_rank_ids.</span>
<span class="sd">        recv_type (type): Data type which received from recv_rank_ids</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with 2 devices. Refer to the tutorial &gt; Distributed Training on mindspore.cn</span>
<span class="sd">        &gt;&gt;&gt; import os</span>
<span class="sd">        &gt;&gt;&gt; import mindspore as ms</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import context</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.neighborexchange = ops.NeighborExchange(send_rank_ids=[1], recv_rank_ids=[1],</span>
<span class="sd">        ...                                                      recv_shapes=([2, 2],), send_shapes=([3, 3],),</span>
<span class="sd">        ...                                                      recv_type=ms.float32)</span>
<span class="sd">        ...</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         out = self.neighborexchange((x,))</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE, device_target=&#39;Ascend&#39;)</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; input_x = Tensor(np.ones([3, 3]), dtype = ms.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[2. 2.], [2. 2.]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">send_rank_ids</span><span class="p">,</span> <span class="n">recv_rank_ids</span><span class="p">,</span> <span class="n">recv_shapes</span><span class="p">,</span> <span class="n">send_shapes</span><span class="p">,</span> <span class="n">recv_type</span><span class="p">,</span>
                 <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_prim_io_names</span><span class="p">(</span><span class="n">inputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;x&#39;</span><span class="p">],</span> <span class="n">outputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;output&#39;</span><span class="p">])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">send_rank_ids</span> <span class="o">=</span> <span class="n">send_rank_ids</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">recv_rank_ids</span> <span class="o">=</span> <span class="n">recv_rank_ids</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">recv_shapes</span> <span class="o">=</span> <span class="n">recv_shapes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">send_shapes</span> <span class="o">=</span> <span class="n">send_shapes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">recv_type</span> <span class="o">=</span> <span class="n">recv_type</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>


<div class="viewcode-block" id="AlltoAll"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.AlltoAll.html#mindspore.ops.AlltoAll">[docs]</a><span class="k">class</span> <span class="nc">AlltoAll</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    AlltoAll is a collective operation.</span>

<span class="sd">    AlltoAll sends data from the all processes to the all processes in the specified group. It has two phases:</span>

<span class="sd">    - The scatter phase: On each process, the operand is split into split_count number of blocks along the</span>
<span class="sd">      split_dimensions, and the blocks are scattered to all processes, e.g., the ith block is send to the ith process.</span>
<span class="sd">    - The gather phase: Each process concatenates the received blocks along the concat_dimension.</span>

<span class="sd">    Note:</span>
<span class="sd">        The tensors must have the same shape and format in all processes of the collection. The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">        This operator requires a full-mesh network topology, each device has the same vlan id, and the ip &amp; mask are</span>
<span class="sd">        in the same subnet, please check the `details \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/programming_guide/zh-CN/r1.6/distributed_training_ops.html#id2&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        split_count (int): On each process, divide blocks into split_count number.</span>
<span class="sd">        split_dim (int): On each process, split blocks along the split_dim.</span>
<span class="sd">        concat_dim (int): On each process, gather the received blocks along the concat_dimension.</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If group is not a string.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with 8 devices. Refer to the tutorial &gt; Distributed Training on mindspore.cn</span>
<span class="sd">        &gt;&gt;&gt; import os</span>
<span class="sd">        &gt;&gt;&gt; import mindspore as ms</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import context</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.alltoall = ops.AlltoAll(split_count = 8, split_dim = -2, concat_dim = -1)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         out = self.alltoall(x)</span>
<span class="sd">        ...         return out</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE, device_target=&#39;Ascend&#39;)</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; rank_id = int(os.getenv(&quot;RANK_ID&quot;))</span>
<span class="sd">        &gt;&gt;&gt; input_x = Tensor(np.ones([1, 1, 8, 1]) * rank_id, dtype = ms.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[[0. 1. 2. 3. 4. 5. 6. 7.]]]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">split_count</span><span class="p">,</span> <span class="n">split_dim</span><span class="p">,</span> <span class="n">concat_dim</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize AlltoAll&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">),</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_is_int</span><span class="p">(</span><span class="n">split_count</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_is_int</span><span class="p">(</span><span class="n">split_dim</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_is_int</span><span class="p">(</span><span class="n">concat_dim</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">split_count</span> <span class="o">=</span> <span class="n">split_count</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">split_dim</span> <span class="o">=</span> <span class="n">split_dim</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">concat_dim</span> <span class="o">=</span> <span class="n">concat_dim</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="n">rank_size</span> <span class="o">=</span> <span class="n">get_group_size</span><span class="p">(</span><span class="n">_get_group</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">group</span><span class="p">))</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">split_count</span> <span class="o">!=</span> <span class="n">rank_size</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;split_count&#39; must be equal to &#39;rank_size&#39;, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;split_count&#39;: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">split_count</span><span class="si">}</span><span class="s2">, &#39;rank_size&#39;: </span><span class="si">{</span><span class="n">rank_size</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">split_dim</span><span class="p">]</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">split_count</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="si">}</span><span class="s2">&#39;, the &#39;split_count&#39; must be divisible by &#39;rank_size&#39;, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;but got &#39;split_count&#39; </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">split_count</span><span class="si">}</span><span class="s2">, &#39;rank_size&#39; </span><span class="si">{</span><span class="n">x_shape</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">split_dim</span><span class="p">]</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
        <span class="n">x_shape</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">concat_dim</span><span class="p">]</span> <span class="o">=</span> <span class="n">x_shape</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">concat_dim</span><span class="p">]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">split_count</span>
        <span class="n">x_shape</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">split_dim</span><span class="p">]</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">x_shape</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">split_dim</span><span class="p">]</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">split_count</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s1">&#39;x&#39;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">target_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">x_dtype</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

<div class="viewcode-block" id="NeighborExchangeV2"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.NeighborExchangeV2.html#mindspore.ops.NeighborExchangeV2">[docs]</a><span class="k">class</span> <span class="nc">NeighborExchangeV2</span><span class="p">(</span><span class="n">Primitive</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    NeighborExchangeV2 is a collective operation.</span>

<span class="sd">    NeighborExchangeV2 sends data from the local rank to ranks in the send_rank_ids,</span>
<span class="sd">    as while receive data from recv_rank_ids.</span>

<span class="sd">    Note:</span>
<span class="sd">        The user needs to preset</span>
<span class="sd">        communication environment variables before running the following example, please check the details on the</span>
<span class="sd">        official website of `MindSpore \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/api/zh-CN/r1.6/api_python/mindspore.ops.html#communication-operators&gt;`_.</span>

<span class="sd">        This operator requires a full-mesh network topology, each device has the same vlan id, and the ip &amp; mask are</span>
<span class="sd">        in the same subnet, please check the `details \</span>
<span class="sd">        &lt;https://www.mindspore.cn/docs/programming_guide/zh-CN/r1.6/distributed_training_ops.html#id2&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        send_rank_ids (list(int)): Ranks which the data is sent to. 8 rank_ids represents 8 directions, if one</span>
<span class="sd">                                   direction is not send to , set it -1.</span>
<span class="sd">        recv_rank_ids (list(int)): Ranks which the data is received from. 8 rank_ids represents 8 directions,</span>
<span class="sd">                                   if one direction is not recv from , set it -1.</span>
<span class="sd">        send_lens (list(int)): Data lens which send to the send_rank_ids, 4 numbers represent the lens of</span>
<span class="sd">                               [top, bottom, left, right].</span>
<span class="sd">        recv_lens (list(int)): Data lens which received from recv_rank_ids, 4 numbers represent the lens of</span>
<span class="sd">                               [top, bottom, left, right].</span>
<span class="sd">        data_format (str): Data format, only support NCHW now.</span>
<span class="sd">        group (str): The communication group to work on. Default: &quot;GlobalComm.WORLD_COMM_GROUP&quot;.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; # This example should be run with 2 devices. Refer to the tutorial &gt; Distributed Training on mindspore.cn</span>
<span class="sd">        &gt;&gt;&gt; import os</span>
<span class="sd">        &gt;&gt;&gt; import mindspore as ms</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import context</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.communication import init</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.neighborexchangev2 = ops.NeighborExchangeV2(send_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1],</span>
<span class="sd">        ...                                                          send_lens=[0, 1, 0, 0],</span>
<span class="sd">        ...                                                          recv_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1],</span>
<span class="sd">        ...                                                          recv_lens=[0, 1, 0, 0],</span>
<span class="sd">        ...                                                          data_format=&quot;NCHW&quot;)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         out = self.neighborexchangev2(x)</span>
<span class="sd">        ...         return out</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE, device_target=&#39;Ascend&#39;)</span>
<span class="sd">        &gt;&gt;&gt; init()</span>
<span class="sd">        &gt;&gt;&gt; input_x = Tensor(np.ones([1, 1, 2, 2]), dtype = ms.float32)</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(input_x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[[1. 1.], [1. 1.], [2. 2.]]]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">send_rank_ids</span><span class="p">,</span> <span class="n">send_lens</span><span class="p">,</span> <span class="n">recv_rank_ids</span><span class="p">,</span> <span class="n">recv_lens</span><span class="p">,</span> <span class="n">data_format</span><span class="p">,</span>
                 <span class="n">group</span><span class="o">=</span><span class="n">GlobalComm</span><span class="o">.</span><span class="n">WORLD_COMM_GROUP</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_prim_io_names</span><span class="p">(</span><span class="n">inputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;x&#39;</span><span class="p">],</span> <span class="n">outputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;output&#39;</span><span class="p">])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">send_rank_ids</span> <span class="o">=</span> <span class="n">send_rank_ids</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">recv_rank_ids</span> <span class="o">=</span> <span class="n">recv_rank_ids</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">send_lens</span> <span class="o">=</span> <span class="n">send_lens</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">recv_lens</span> <span class="o">=</span> <span class="n">recv_lens</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">format</span> <span class="o">=</span> <span class="n">data_format</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;group&#39;</span><span class="p">,</span> <span class="n">_get_group</span><span class="p">(</span><span class="n">group</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;no_eliminate&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

<span class="k">class</span> <span class="nc">_MirrorOperator</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do all reduce and mean in backward. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group to work on. Default: None.</span>
<span class="sd">        dev_num (int): The device number of the group. Default: None.</span>
<span class="sd">        mean_flag (bool): Whether use mean in backward. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dev_num</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">mean_flag</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _MirrorOperator.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">group</span> <span class="o">=</span> <span class="n">group</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dev_num</span> <span class="o">=</span> <span class="n">dev_num</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mean_flag</span> <span class="o">=</span> <span class="n">mean_flag</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;fusion&quot;</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="n">mirror</span> <span class="o">=</span> <span class="n">_MirrorOperator</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">_MirrorMiniStepOperator</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do all reduce and mean in backward. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group to work on. Default: None.</span>
<span class="sd">        dev_num (int): The device number of the group. Default: None.</span>
<span class="sd">        mean_flag (bool): Whether use mean in backward. Default: None.</span>
<span class="sd">        grad_accumulation_step (int): The grad accumulation step. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dev_num</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">mean_flag</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">grad_accumulation_step</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _MirrorMiniStepOperator.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">group</span> <span class="o">=</span> <span class="n">group</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dev_num</span> <span class="o">=</span> <span class="n">dev_num</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mean_flag</span> <span class="o">=</span> <span class="n">mean_flag</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">grad_accumulation_step</span> <span class="o">=</span> <span class="n">grad_accumulation_step</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="n">mirror_mini_step</span> <span class="o">=</span> <span class="n">_MirrorMiniStepOperator</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">_VirtualDiv</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do Div in backward.</span>

<span class="sd">    Args:</span>
<span class="sd">        divisor: float32</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">divisor</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _VirtualDiv.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">divisor</span> <span class="o">=</span> <span class="n">divisor</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="n">virtual_div</span> <span class="o">=</span> <span class="n">_VirtualDiv</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">_VirtualAdd</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Auto parallel virtual operator. Do nothing in forward, do Add in backward.&quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _VirtualAdd.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">y_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">y_dtype</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">_VirtualDataset</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual dataset operator.</span>

<span class="sd">    It would insert VirtualDataset operator in forward computation and be deleted before backward computation.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _VirtualDataset.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">args</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">args</span>


<span class="n">virtual_dataset</span> <span class="o">=</span> <span class="n">_VirtualDataset</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">_VirtualAssignAdd</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do AssignAdd in backward. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _VirtualAssignAdd.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">y_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">y_dtype</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="n">virtual_assign_add</span> <span class="o">=</span> <span class="n">_VirtualAssignAdd</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">_VirtualAccuGrad</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, return y in backward. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _VirtualAccuGrad.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">y_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">y_dtype</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="n">virtual_accu_grad</span> <span class="o">=</span> <span class="n">_VirtualAccuGrad</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">_MirrorMicroStepOperator</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual operator. Do nothing in forward, do all reduce and mean in backward. It is only for</span>
<span class="sd">    internal use of parallel modules and cannot be called by users.</span>

<span class="sd">    Args:</span>
<span class="sd">        group (str): The communication group to work on. Default: None.</span>
<span class="sd">        dev_num (int): The device number of the group. Default: None.</span>
<span class="sd">        mean_flag (bool): Whether use mean in backward. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dev_num</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">mean_flag</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _MirrorMicroStepOperator.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">group</span> <span class="o">=</span> <span class="n">group</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dev_num</span> <span class="o">=</span> <span class="n">dev_num</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mean_flag</span> <span class="o">=</span> <span class="n">mean_flag</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">z_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">_VirtualOutput</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Auto parallel virtual out operator.</span>

<span class="sd">    It would insert VirtualOutput operator in forward computation and be deleted before backward computation.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _VirtualOutput.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x_dtype</span>


<span class="k">class</span> <span class="nc">_GetTensorSlice</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gets tensor slice by device matrix and tensor map.</span>

<span class="sd">    Args:</span>
<span class="sd">        dev_mat (tuple): The device matrix of the slice tensor.</span>
<span class="sd">        tensor_map (tuple): The tensor map of the slice tensor.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _GetTensorSlice.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">infer_value</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">dev_mat</span><span class="p">,</span> <span class="n">tensor_map</span><span class="p">):</span>
        <span class="kn">from</span> <span class="nn">mindspore.parallel._tensor</span> <span class="kn">import</span> <span class="n">_load_tensor</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;dev_mat&quot;</span><span class="p">,</span> <span class="n">dev_mat</span><span class="p">,</span> <span class="p">[</span><span class="nb">tuple</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;tensor_map&quot;</span><span class="p">,</span> <span class="n">tensor_map</span><span class="p">,</span> <span class="p">[</span><span class="nb">tuple</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">Tensor</span><span class="p">(</span><span class="n">_load_tensor</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">dev_mat</span><span class="p">,</span> <span class="n">tensor_map</span><span class="p">))</span>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>